
import db
from datetime import datetime, timedelta
from decorators import cached_property
from utils.format import d2id
from controllers.trade_cal import TradeCalController
from collections import defaultdict

class StrategyController(object):

    def __init__(self, id=None, user=None):
        self.id = id
        self.user = user

    @cached_property
    def strategy(self):
        return db.StrategyTB.get(id=self.id)

    @classmethod
    def all_strategys(cls):
        """获取所有指标"""
        return db.StrategyTB.find()

    @classmethod
    def create(cls, **args):
        db.StrategyTB.add(**args)

    def delete(self):
        db.StrategyTB.rm(id=self.id)

    def update(self, **args):
        db.StrategyTB.set(**args).where(id=self.id)

    def get_strategy_trigger_stock(self, cal_date_int=None):
        """取某个策略下所有触发的股票"""
        next_trade_day = db.TradeCalTB.next_trade_day(cal_date_int)

        records = db.TriggerRecordTB.select().where(
                strategy=self.strategy["name"], trigger_date=next_trade_day
                ).order_by("orient").all()
        all_ts_code = [i["ts_code"] for i in records]
        all_trigger_stock = db.StockTB.find(ts_code=("in", all_ts_code))
        stock_name_map = {i["ts_code"]:i["name"] for i in all_trigger_stock}
        for i in records:
            i['name'] = stock_name_map.get(i["ts_code"]) or ''
        return records

    @classmethod
    def strategy_statistic(cls, cal_date_int=None):
        """对所有今日触发的股票做统计"""
        # 买入触发个数
        next_trade_day = db.TradeCalTB.next_trade_day(cal_date_int)
        all_trigger = db.TriggerRecordTB.find(trigger_date=next_trade_day)
        all_strategys = db.StrategyTB.find(available=1)
        trigger_statistic = {i["name"]:{
                    "id":i["id"],
                    "strategy":i["name"],
                    "orient":i["orient"],
                    "count": 0
                    } for i in all_strategys}
        for i in all_trigger:
            if trigger_statistic.get(i["strategy"]):
                trigger_statistic[i["strategy"]]["count"] += 1
        return {"data": list(trigger_statistic.values()), "trade_date": next_trade_day}


    @classmethod
    def trigger_statistic(cls, cal_date_int):
        next_trade_day = db.TradeCalTB.next_trade_day(cal_date_int)
        all_trigger = db.TriggerRecordTB.find(trigger_date=next_trade_day)
        buy_stock = defaultdict(int)
        sell_stock = defaultdict(int)
        for i in all_trigger:
            if i["orient"] == 1:
                buy_stock[i["ts_code"]] += 1
            else:
                sell_stock[i["ts_code"]] += 1
        all_ts_code = list(buy_stock.keys()) + list(sell_stock.keys())
        all_trigger_stock = db.StockTB.find(ts_code=("in", all_ts_code))
        stock_name_map = {i["ts_code"]:i["ts_code"] + ' ' +i["name"] for i in all_trigger_stock}
        ts_code_name_buy = list()
        ts_code_name_sell = list()
        for k, v in buy_stock.items():
            ts_code_name_buy.append({"stock": stock_name_map[k], "count": v, "ts_code":k})
        ts_code_name_buy.sort(key=lambda x:x["count"], reverse=True)
        for k, v in sell_stock.items():
            ts_code_name_sell.append({"stock": stock_name_map[k], "count": v, "ts_code":k})
        ts_code_name_buy.sort(key=lambda x:x["count"], reverse=True)

        return {"buy": ts_code_name_buy, "sell": ts_code_name_sell, "trade_date":next_trade_day}







